#pragma once 
#include <set>
#include "../common/net.hpp"
#include "../common/message.hpp"

namespace yjz_rpc
{
    namespace server
    {
        class ProviderManager
        {
        public:
            using ptr = std::shared_ptr<ProviderManager>;
            struct Provider
            {
                using ptr = std::shared_ptr<Provider>;
                Provider(const BaseConnection::ptr& c, const Address& h) // 构造函数
                    : conn(c), host(h)
                {}

                void appendMethod(const std::string& method) // 增加服务
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    methods.emplace_back(method);
                }

                Address host; // 外部可访问地址
                std::mutex mutex;
                BaseConnection::ptr conn; // provider对应的连接
                std::vector<std::string> methods; // 当前provider所能提供的服务
            };

            // 有新的服务提供者进行注册的时候，新增provider
            void addProvider(const BaseConnection::ptr& conn, const Address& host, const std::string& method)
            {
                Provider::ptr provider;
                // 1.查找当前provider是否已经存在
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it = _conns.find(conn);
                    if (it != _conns.end())
                    {
                        provider = it->second;
                    }
                    else
                    {
                        provider = std::make_shared<Provider>(conn, host);
                        _conns.insert(std::make_pair(conn, provider));
                    }
                    // 2.不管这个服务是否已经有人提供，这个服务的提供者都多了一个
                    _providers[method].insert(provider);
                }
                // 3.向服务对象中新增一个所能提供的服务名称
                provider->appendMethod(method);
            }

            // 有服务提供者下线，查找provider，用于下线通知
            Provider::ptr getProvider(const BaseConnection::ptr& conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conns.find(conn);
                if (it == _conns.end())
                {
                    return Provider::ptr();
                }
                return it->second;
            }

            // 有服务提供者下线，删除他的所有关联信息
            void delProvider(const BaseConnection::ptr& conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conns.find(conn);
                if (it == _conns.end())
                {
                    return;
                }
                for (auto& method : it->second->methods)
                {
                    _providers[method].erase(it->second);
                }
                _conns.erase(it);
            }

            std::vector<Address> methodHosts(const std::string& method)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _providers.find(method);
                if (it == _providers.end())
                {
                    LOG(LogLevel::WARNING) << "没有主机提供该服务！";
                    return std::vector<Address>();
                }
                std::vector<Address> result;
                for (auto& provider : it->second)
                {
                    result.push_back(provider->host);
                }
                return result;
            }
        private:
            std::mutex _mutex;
            std::unordered_map<BaseConnection::ptr, Provider::ptr> _conns; // provider对应的连接
            std::unordered_map<std::string, std::set<Provider::ptr>> _providers; // 一个服务可以由哪些provider提供
        };

        class DiscovererManager
        {
        public:
            using ptr = std::shared_ptr<DiscovererManager>;
            struct Discoverer
            {
                using ptr = std::shared_ptr<Discoverer>;
                Discoverer(const BaseConnection::ptr& c)
                    : conn(c)
                {}

                void appendMethod(const std::string& method)
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    methods.emplace_back(method);
                }

                std::mutex mutex;
                BaseConnection::ptr conn; // 发现者关联的客户端连接
                std::vector<std::string> methods; // 该服务发现者发现过的服务名称
            };

            // 当客户端进行服务发现时，新增discoverer
            Discoverer::ptr addDiscoverer(const BaseConnection::ptr& conn, const std::string& method)
            {
                Discoverer::ptr discoverer;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it = _conns.find(conn);
                    if (it != _conns.end())
                    {
                        discoverer = it->second;
                    }
                    else
                    {
                        discoverer = std::make_shared<Discoverer>(conn);
                        _conns.insert(std::make_pair(conn, discoverer));
                    }
                    // 这个服务的发现者+1
                    _discoverers[method].insert(discoverer);
                }
                discoverer->appendMethod(method);
                return discoverer;
            }

            // 当客户端断开连接时，删除discoverer信息
            void delDiscoverer(const BaseConnection::ptr& conn)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _conns.find(conn);
                if (it == _conns.end())
                {
                    return;
                }
                else
                {
                    // discoverer发现过哪些服务，删除这些信息
                    for (auto& method : it->second->methods)
                    {
                        _discoverers[method].erase(it->second);
                    }
                }
                _conns.erase(it);
            }

            // 当有新的服务提供者上线，进行上线通知
            void onlineNotify(const std::string& method, const Address& host)
            {
                Notify(method, host, ServiceOpType::SERVICE_ONLINE);
            }

            // 下线通知
            void offlineNotify(const std::string& method, const Address& host)
            {
                Notify(method, host, ServiceOpType::SERVICE_OFFLINE);
            }
        private:
            // 服务上线/下线，都有哪些客户发现过这个服务，进行推送
            void Notify(const std::string& method, const Address& host, ServiceOpType type)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                auto it = _discoverers.find(method);
                if (it == _discoverers.end())
                {
                    return;
                }
                auto req_msg = MessageFactory::create<ServiceRequest>();
                req_msg->setHost(host);
                req_msg->setId(UUID::uuid());
                req_msg->setMethod(method);
                req_msg->setMType(MType::REQ_SERVICE);
                req_msg->setOpType(type);
                for (auto& discover : it->second)
                {
                    discover->conn->send(req_msg);
                }
            }
        private:
            std::mutex _mutex;
            std::unordered_map<BaseConnection::ptr, Discoverer::ptr> _conns; // discoverer对应的连接
            std::unordered_map<std::string, std::set<Discoverer::ptr>> _discoverers; // 哪些服务被哪些主机发现过
        };

        class ProDisManager
        {
        public:
            using ptr = std::shared_ptr<ProDisManager>;
            ProDisManager()
                : _providers(std::make_shared<ProviderManager>())
                , _discoverers(std::make_shared<DiscovererManager>())
            {}

            // 提供给Dispatcher模块，针对不同的请求做处理
            void onServiceRequest(const BaseConnection::ptr& conn, const ServiceRequest::ptr& ser_req)
            {
                ServiceOpType type = ser_req->opType();
                // 服务注册
                if (type == ServiceOpType::SERVICE_REGISTRY)
                {
                    // 新增服务提供者，然后进行上线通知
                    LOG(LogLevel::INFO) << ser_req->host().first << " : " << ser_req->host().second << " 注册了服务： " << ser_req->method();
                    _providers->addProvider(conn, ser_req->host(), ser_req->method());
                    _discoverers->onlineNotify(ser_req->method(), ser_req->host());
                    // 进行响应
                    registryResponse(conn, ser_req);
                }
                // 服务发现
                else if (type == ServiceOpType::SERVICE_DISCOVERY)
                {
                    LOG(LogLevel::INFO) << "客户端进行 " << ser_req->method() << " 服务发现";
                    _discoverers->addDiscoverer(conn, ser_req->method());
                    discoveryResponse(conn, ser_req);
                }
                else
                {
                    LOG(LogLevel::ERROR) << "服务请求类型错误！";
                    errorResponse(conn, ser_req);
                }
            }
            
            // 连接关闭的回调
            void onConnShutDown(const BaseConnection::ptr& conn)
            {
                auto provider = _providers->getProvider(conn);
                if (provider.get() != nullptr)
                {
                    LOG(LogLevel::INFO) << provider->host.first << " : " << provider->host.second << " 服务下线!";
                    // 当服务提供者下线时，对他所提供的服务进行下线通知
                    for (auto& method : provider->methods)
                    {
                        _discoverers->offlineNotify(method, provider->host);
                    }
                    _providers->delProvider(conn);
                }
                // 直接删除不需要判断，因为del接口中会查找，没有的话会自己退出
                _discoverers->delDiscoverer(conn);
            }
        private:
            void errorResponse(const BaseConnection::ptr& conn, const ServiceRequest::ptr& ser_req)
            {
                auto rsp_msg = MessageFactory::create<ServiceResponse>();
                rsp_msg->setId(ser_req->getId());
                rsp_msg->setMType(MType::RSP_SERVICE);
                rsp_msg->setOpType(ServiceOpType::SERVICE_UNKNOW);
                rsp_msg->setRcode(RCode::RCODE_INVALID_OPTYPE);
                conn->send(rsp_msg);
            }

            void registryResponse(const BaseConnection::ptr& conn, const ServiceRequest::ptr& ser_req)
            {
                auto rsp_msg = MessageFactory::create<ServiceResponse>();
                rsp_msg->setId(ser_req->getId());
                rsp_msg->setMType(MType::RSP_SERVICE);
                rsp_msg->setOpType(ServiceOpType::SERVICE_REGISTRY);
                rsp_msg->setRcode(RCode::RCODE_OK);
                conn->send(rsp_msg);
            }

            void discoveryResponse(const BaseConnection::ptr& conn, const ServiceRequest::ptr& ser_req)
            {
                auto rsp_msg = MessageFactory::create<ServiceResponse>();
                rsp_msg->setId(ser_req->getId());
                rsp_msg->setMType(MType::RSP_SERVICE);
                rsp_msg->setOpType(ServiceOpType::SERVICE_DISCOVERY);
                std::vector<Address> hosts(_providers->methodHosts(ser_req->method()));
                if (hosts.empty())
                {
                    rsp_msg->setRcode(RCode::RCODE_NOT_FOUND_SERVICE);
                }
                else
                {
                    rsp_msg->setRcode(RCode::RCODE_OK);
                    rsp_msg->setHosts(hosts);
                    rsp_msg->setMethod(ser_req->method());
                }
                conn->send(rsp_msg);
            }
        private:
            ProviderManager::ptr _providers;
            DiscovererManager::ptr _discoverers;
        };
    }
}